import subprocess
import multiprocessing as mp
from multiprocessing import Process
import re
import logging
import os
import sys
import asyncio
from fastapi import FastAPI
import uvicorn
from flask import Flask
from fastchat.utils import StreamToLogger
from PyCmpltrtok.common import get_dir_name_ext
from python_nlp.xfaiss.flask_faiss_on_mongo import logger as logger_faiss
from python_nlp.xfaiss.flask_faiss_on_mongo_fastapi import logger as logger_faiss_fastapi
from python_nlp.kg.neo4j.load_data.data2mongo_words import EMBED_TBL as EMBED_TBL_KG_WORDS

NEO4J_NAME = 'neo4j-agr'

logger = logging.getLogger(__name__)

LOG_FORMAT = "%(asctime)s - %(pathname)s[line:%(lineno)d] - %(levelname)s: %(name)s: %(message)s"
DIR, BASE, EXT = get_dir_name_ext(os.path.abspath(__file__))
PATH = os.path.join(DIR, 'log.tmp.d', f'{BASE}.tmp.log')

logging.basicConfig(
    level=logging.DEBUG,
    format=LOG_FORMAT,
)

logger_faiss.setLevel(logging.DEBUG)
logger_faiss_fastapi.setLevel(logging.DEBUG)

# hdl_out = logging.StreamHandler()
# hdl_out.setLevel(logging.DEBUG)
# hdl_out.setFormatter(logging.Formatter(LOG_FORMAT))

# logger.addHandler(hdl_out)
# logger_faiss.addHandler(hdl_out)
# logger_faiss_fastapi.addHandler(hdl_out)

# hdl = logging.FileHandler(filename=PATH)
hdl = logging.handlers.TimedRotatingFileHandler(
    PATH, when="M", interval=10, utc=True, encoding="utf-8"
)
hdl.setLevel(logging.DEBUG)
hdl.setFormatter(logging.Formatter(LOG_FORMAT))

logger.addHandler(hdl)
logger_faiss.addHandler(hdl)
logger_faiss_fastapi.addHandler(hdl)


def std2log():
    # Redirect stdout and stderr to loggers
    stdout_logger = logging.getLogger("stdout")
    stdout_logger.setLevel(logging.INFO)
    stdout_logger.addHandler(hdl)
    # stdout_logger.addHandler(hdl_out)
    sl = StreamToLogger(stdout_logger, logging.INFO)
    sys.stdout = sl

    stderr_logger = logging.getLogger("stderr")
    stderr_logger.setLevel(logging.ERROR)
    stderr_logger.addHandler(hdl)
    # stderr_logger.addHandler(hdl_out)
    sl = StreamToLogger(stderr_logger, logging.ERROR)
    sys.stderr = sl
    

def get_neo4j_docker_port(target_name):
    output = subprocess.check_output(["docker", "ps", "-a"]).decode('utf8')
    logger.debug('docker ps -a => %s', output)
    lines = output.split('\n')
    lines = lines[1:]
    for line in lines:
        port, id = parse_neo4j_docker_ps(line, target_name)
        if port is None:
            continue
        return port, id
    return None, None


parse_neo4j_docker_ps_regexp = re.compile(r'0\.0\.0\.0\:(\d+)\-\>7474\/tcp')


def parse_neo4j_docker_ps(line, target_name):
    columns = line.split(' ' * 3)
    columns = [c.strip() for c in columns]
    if columns[-1] != target_name:
        return None, None
    ports_str = columns[-2]
    match = parse_neo4j_docker_ps_regexp.search(ports_str)
    if not match:
        logger.info('%s is not running', target_name)
        return 0, columns[0]
    port = int(match[1])
    logger.info('neo4j is on port %d', port)
    return port, columns[0]


if 0:            
    def _set_app_event_flask(app: Flask, started_event: mp.Event = None):
        """
        WWW: @app.on_event("startup") equivalent in Flask
        
        @app.before_request is not desirable: it is run just before the 1st request.
        """
        @app.before_request
        def on_startup():
            if started_event is not None:
                print('Event set.', flush=True)
                started_event.set()
                

    def run_faiss_search_service(port, model_name='moka-ai/m3e-base', dev=0,
        host='0.0.0.0', mongo_link='local', mongo_table=EMBED_TBL_KG_WORDS, 
        started_event: mp.Event = None):
        std2log()
        
        from python_nlp.xfaiss.flask_faiss_on_mongo import app as app_faiss
        from python_nlp.xfaiss.flask_faiss_on_mongo import (
            load_embed_model,
            build_search_index
        )
        
        dev, model, tokenizer, max_len = load_embed_model(dev, model_name)
        global_cache_map = dict()
        index, xsentences, xmd5s = build_search_index(mongo_link, mongo_table)
        sys.modules["python_nlp.xfaiss.flask_faiss_on_mongo"].dev = dev
        sys.modules["python_nlp.xfaiss.flask_faiss_on_mongo"].model = model
        sys.modules["python_nlp.xfaiss.flask_faiss_on_mongo"].tokenizer = tokenizer
        sys.modules["python_nlp.xfaiss.flask_faiss_on_mongo"].max_len = max_len
        sys.modules["python_nlp.xfaiss.flask_faiss_on_mongo"].global_cache_map = global_cache_map
        sys.modules["python_nlp.xfaiss.flask_faiss_on_mongo"].index = index
        sys.modules["python_nlp.xfaiss.flask_faiss_on_mongo"].xsentences = xsentences
        sys.modules["python_nlp.xfaiss.flask_faiss_on_mongo"].xmd5s = xmd5s
        
        _set_app_event_flask(app_faiss, started_event)
        app_faiss.run(host=host, port=port)


def _set_app_event(app: FastAPI, started_event: mp.Event = None):
    @app.on_event("startup")
    async def on_startup():
        if started_event is not None:
            print('Event set.', flush=True)
            started_event.set()  
            
            
def run_faiss_search_service_fastapi(port, model_name='moka-ai/m3e-base', dev=0,
    host='0.0.0.0', mongo_link='local', mongo_table=EMBED_TBL_KG_WORDS, 
    started_event: mp.Event = None):
    std2log()
    
    from python_nlp.xfaiss.flask_faiss_on_mongo_fastapi import app as app_faiss
    from python_nlp.xfaiss.flask_faiss_on_mongo_fastapi import (
        load_embed_model,
        build_search_index
    )
    
    dev, model, tokenizer, max_len = load_embed_model(dev, model_name)
    global_cache_map = dict()
    index, xsentences, xmd5s = build_search_index(mongo_link, mongo_table)
    sys.modules["python_nlp.xfaiss.flask_faiss_on_mongo_fastapi"].dev = dev
    sys.modules["python_nlp.xfaiss.flask_faiss_on_mongo_fastapi"].model = model
    sys.modules["python_nlp.xfaiss.flask_faiss_on_mongo_fastapi"].tokenizer = tokenizer
    sys.modules["python_nlp.xfaiss.flask_faiss_on_mongo_fastapi"].max_len = max_len
    sys.modules["python_nlp.xfaiss.flask_faiss_on_mongo_fastapi"].global_cache_map = global_cache_map
    sys.modules["python_nlp.xfaiss.flask_faiss_on_mongo_fastapi"].index = index
    sys.modules["python_nlp.xfaiss.flask_faiss_on_mongo_fastapi"].xsentences = xsentences
    sys.modules["python_nlp.xfaiss.flask_faiss_on_mongo_fastapi"].xmd5s = xmd5s
    
    _set_app_event(app_faiss, started_event)
    uvicorn.run(app_faiss, host=host, port=port)
            
            
if '__main__' == __name__:
    
    async def _main():
        std2log()
        
        mp.set_start_method("spawn")
        manager = mp.Manager()
        name2tuple = {}
        
        #############################################################################################################
        # Start neo4j docker
        port, id = get_neo4j_docker_port(NEO4J_NAME)
        print('port:', port)
        
        need_rm, need_start = False, False
        if port is None:
            need_start = True
        elif 0 == port:
            need_rm = True
            need_start = True
        
        logger.info(f'need_rm: {need_rm}, need_start: {need_start}')
        if need_rm:
            output = subprocess.check_output(["docker", "rm", id]).decode('utf8')
            logger.info(f'rm existed neo4j container {id}: {output}')
        if need_start:
            cmd = """docker run --name neo4j-agr -d \
-p 7681:7687 -p 7468:7474 \
-e NEO4J_apoc_export_file_enabled=true \
-e NEO4J_apoc_import_file_enabled=true \
-e NEO4J_apoc_import_file_use__neo4j__config=true \
-e NEO4JLABS_PLUGINS=["apoc"] \
-v neo4j-agr_tmp:/tmp \
-v neo4j-agr_plugins:/plugins \
-v neo4j-agr:/data \
neo4j:5.11.0-community"""
            cmd_arr = cmd.split(' ')
            output = subprocess.check_output(cmd_arr).decode('utf8')
            logger.info(f'Start neo4j container: {cmd_arr}: {output}')
    
        #############################################################################################################
        # Prepare to start FAISS search service to seek node in neo4j
        faiss_service_started = manager.Event()
        process = Process(
            target=run_faiss_search_service_fastapi,
            name="faiss_search_service",
            kwargs=dict(
                port=7700,
                model_name='/home/yunpeng/models/hf/m3e-base',
                dev=1,
                started_event=faiss_service_started,
            ),
            daemon=False,
        )
        name2tuple['faiss_search'] = (process, faiss_service_started)
        
        #############################################################################################################
        # Run all the processes
        for k, (p, e) in name2tuple.items():
            logger.info(f'Starting {k} ...')
            p.start()
            p.name = f"{p.name} ({p.pid})"
            e.wait()
            logger.info(f'Starting {k} Done')
            
        while True:
            print('Half a minute is passed.', file=sys.stderr, flush=True)
            print('Half a minute is passed.', flush=True)
            await asyncio.sleep(30.0)
    
    asyncio.run(_main())
        